
import os
import numpy as np
import matplotlib as plt
from scipy.interpolate import griddata
from simple_tif import array_to_tif, tif_to_array
from parameters import (direct_lattice_vectors,
                        corrected_shift_vector,
                        offset_vector)
from array_illumination import generate_lattice, get_shift

###########load data##############

def read_data(file_name, z, x, y):
    return np.fromfile(file_name, dtype = np.uint16
                   ).reshape((z,x,y)).astype(np.float32)
zpix=400
filenames = ['nl_image_Z_%05inm.raw'%i for i in range(44595, 49996, 200)]
for filename in filenames:
    print "Processing", filename
    basename, ext = os.path.splitext(filename)

    if ext == '.raw':
        data = read_data(filename, zpix, 256, 256)
    elif ext == '.tif':
        data = tif_to_array(filename).astype(np.float32)
    print data.shape

    points = []
    values = []
    pinhole_radius = 20
    mask_sigma = 5
    grid_x, grid_y = np.mgrid[
        0:data.shape[1]:8j*data.shape[1],
        0:data.shape[2]:8j*data.shape[2]]
    widefield_image = np.zeros(data.shape, dtype=np.float64)
    print "zpix:",
    for z in range(zpix):
        print z,
        lattice = generate_lattice(
                image_shape=data.shape[1:],
                lattice_vectors=direct_lattice_vectors,
                center_pix=offset_vector + get_shift(
                    corrected_shift_vector, z),
                edge_buffer=pinhole_radius + 1)
        normalization = 0
        num_added = 0
        widefield_contribution = []
        new_values = []
        slices = []
        for x, y in lattice:
            points.append((x, y))
            x_slice = slice(round(x-pinhole_radius), round(x+pinhole_radius+1))
            y_slice = slice(round(y-pinhole_radius), round(y+pinhole_radius+1))
            slices.append((x_slice, y_slice))
            x_coords = np.arange(x_slice.start - x, x_slice.stop - x, 1)
            y_coords = np.arange(y_slice.start - y, y_slice.stop - y, 1)
            x_mask = np.exp(-x_coords**2 / (2*mask_sigma**2)
                            ).reshape(x_coords.size, 1)
            y_mask = np.exp(-y_coords**2 / (2*mask_sigma**2)
                            ).reshape(1, y_coords.size)
            widefield_contribution.append(
                data[z, x_slice, y_slice] * x_mask * y_mask)
            new_values.append(widefield_contribution[-1].sum())
            normalization += new_values[-1]
            num_added += 1
        normalization = normalization * 1.0 / num_added
        for i, (x, y) in enumerate(lattice):
            """
            Now that we have our normalization factor, normalize the
            values we'll use for the excitation image, and add up the
            contributions to the widefield image.
            """
            new_values[i] = new_values[i] * 1.0 / normalization
            widefield_image[z, slices[i][0], slices[i][1]] += (
                widefield_contribution[i] * 1.0 / normalization)
        values.extend(new_values)
    print
    print "Saving preprocessed emission..."
    array_to_tif(widefield_image.reshape(
        widefield_image.shape[0], widefield_image.shape[1], widefield_image.shape[2]
        ).astype(np.float32),
                 basename + '_preprocessed_emission.tif')
    widefield_image = widefield_image.sum(axis=0)
    print "Interpolating excitation..."
    interp_method = 'cubic'
    grid_z2 = griddata(
        np.asarray(points), np.asarray(values),
        (grid_x, grid_y), method=interp_method)
    print "Interpolating emission..."
    widefield_points = np.array([(x, y)
                                 for x in range(widefield_image.shape[0])
                                 for y in range(widefield_image.shape[1])])
    widefield_values = widefield_image.ravel()
    interp_widefield_image = griddata(
        widefield_points, widefield_values,
        (grid_x, grid_y), method=interp_method)

                                
    print "Done interpolating."

    print "Saving..."
    array_to_tif(grid_z2.reshape(1, grid_z2.shape[0], grid_z2.shape[1]
                                 ).astype(np.float32),
                 basename + '_excitation.tif')
    array_to_tif(interp_widefield_image.reshape(
        (1, interp_widefield_image.shape[0], interp_widefield_image.shape[1])
        ).astype(np.float32),
                 basename + '_emission.tif')
    print "Done saving."
